from vnpy_ctastrategy import (
    CtaTemplate,
    StopOrder,
    TickData,
    BarData,
    TradeData,
    OrderData
)
from vnpy.trader.constant import Direction, Offset
import numpy as np

class GridStrategy(CtaTemplate):
    """网格交易策略"""
    author = "Youpeng Hu"

    # 策略参数
    grid_upper = 100.0        # 网格上限
    grid_lower = 50.0        # 网格下限
    grid_number = 10         # 网格数量
    order_volume = 1         # 每单数量

    # 策略变量
    current_tick = None
    active_orders = []       # 活跃订单列表
    grid_levels = []         # 网格价格层级
    position_prices = []     # 持仓成本价列表

    parameters = ["grid_upper", "grid_lower", "grid_number", "order_volume"]
    variables = ["grid_upper", "grid_lower", "grid_number", "order_volume"]

    def __init__(self, cta_engine, strategy_name, vt_symbol, setting):
        super().__init__(cta_engine, strategy_name, vt_symbol, setting)
        
        # 初始化网格价格
        self.generate_grid_levels()
        
    def generate_grid_levels(self):
        """生成网格价格层级"""
        price_range = self.grid_upper - self.grid_lower
        step = price_range / self.grid_number
        self.grid_levels = [round(self.grid_lower + i*step, 2) for i in range(self.grid_number+1)]
        
    def on_init(self):
        """初始化回调"""
        self.write_log("策略初始化")
        self.load_bar(1)  # 加载1分钟数据

    def on_start(self):
        """启动回调"""
        self.write_log("策略启动")
        self.put_event()

    def on_stop(self):
        """停止回调"""
        self.write_log("策略停止")
        self.put_event()

    def on_tick(self, tick: TickData):
        """Tick更新回调"""
        self.current_tick = tick
        
        # 检查并更新网格订单
        self.check_grid_orders()
        
        self.put_event()

    def on_bar(self, bar: BarData):
        """K线更新回调"""
        # 清除已失效的订单
        self.cancel_expired_orders()
        
        # 生成网格订单
        self.generate_grid_orders()
        
        self.put_event()

    def check_grid_orders(self):
        """检查并处理成交订单"""
        # 获取当前价格
        current_price = self.current_tick.last_price
        
        # 检查卖出条件
        if self.pos > 0:
            for price in self.position_prices:
                if current_price >= price * 1.01:  # 当价格上涨1%时卖出
                    self.sell(price * 1.01, self.order_volume)
                    self.position_prices.remove(price)
                    break

        # 检查买入条件
        if len(self.position_prices) < len(self.grid_levels):
            for level in reversed(self.grid_levels):
                if current_price <= level and level not in self.position_prices:
                    self.buy(level, self.order_volume)
                    self.position_prices.append(level)
                    break

    def generate_grid_orders(self):
        """生成网格订单"""
        # 撤消所有未完成订单
        self.cancel_all()
        
        # 生成买入订单
        for price in self.grid_levels:
            if price < self.current_tick.last_price and price not in self.position_prices:
                self.buy(price, self.order_volume)
                
        # 生成卖出订单
        for price in self.position_prices:
            self.sell(price * 1.01, self.order_volume)

    def cancel_expired_orders(self):
        """取消过期订单"""
        for order in self.active_orders:
            if order.status in [OrderStatus.SUBMITTING, OrderStatus.NOTTRADED]:
                self.cancel_order(order.vt_orderid)

    def on_order(self, order: OrderData):
        """订单更新回调"""
        if order.status == OrderStatus.ALLTRADED:
            if order.direction == Direction.LONG:
                self.write_log(f"买入成交：{order.price}")
            else:
                self.write_log(f"卖出成交：{order.price}")
                
        # 更新活跃订单列表
        if order.status == OrderStatus.SUBMITTING:
            self.active_orders.append(order)
        elif order.status in [OrderStatus.CANCELLED, OrderStatus.REJECTED]:
            if order in self.active_orders:
                self.active_orders.remove(order)

    def on_trade(self, trade: TradeData):
        """成交更新回调"""
        self.put_event()

    def buy(self, price, volume):
        """发送买入委托"""
        return self.send_order(Direction.LONG, Offset.OPEN, price, volume)

    def sell(self, price, volume):
        """发送卖出委托"""
        return self.send_order(Direction.SHORT, Offset.CLOSE, price, volume)

    def send_order(self, direction, offset, price, volume):
        """发送委托"""
        return self.buy(price, volume) if direction == Direction.LONG else self.sell(price, volume)